import dagre from 'dagre'
import { Node, Edge } from '@vue-flow/core'

export interface DagreLayoutOptions {
  rankdir?: 'TB' | 'BT' | 'LR' | 'RL' // Direction: Top-Bottom, Bottom-Top, Left-Right, Right-Left
  nodesep?: number // Node spacing
  ranksep?: number // Layer spacing
  edgesep?: number // Edge spacing
}

export class DagreLayoutUtil {
  private graph: dagre.graphlib.Graph

  constructor(options: DagreLayoutOptions = {}) {
    this.graph = new dagre.graphlib.Graph()
    this.graph.setGraph({
      rankdir: options.rankdir || 'TB', // Default: Top-Bottom
      nodesep: options.nodesep || 50,
      ranksep: options.ranksep || 100,
      edgesep: options.edgesep || 10,
    })
    this.graph.setDefaultEdgeLabel(() => ({}))
  }

  applyLayout(nodes: Node[], edges: Edge[]): Node[] {
    nodes.forEach((node) => {
      const width = node.width || 150
      const height = node.height || 50
      this.graph.setNode(node.id, { width, height })
    })

    edges.forEach((edge) => {
      this.graph.setEdge(edge.source, edge.target)
    })

    dagre.layout(this.graph)

    return nodes.map((node) => {
      const dagreNode = this.graph.node(node.id)
      return {
        ...node,
        position: {
          x: dagreNode.x - (node.width || 150) / 2,
          y: dagreNode.y - (node.height || 50) / 2,
        },
      }
    })
  }
}